{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1809c79c",
   "metadata": {},
   "source": [
    "# TVM Relay 自动量化校准\n",
    "\n",
    "参考：`tvm/src/relay/quantize/calibrate.cc` 和 `tvm/python/tvm/relay/quantize/_calibrate.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b8421e64",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "155b85e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "from tvm.relay import analysis as _analysis\n",
    "from tvm.relay import op as _op\n",
    "from tvm.relay.op import op as _reg\n",
    "from tvm.relay import expr as _expr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ead74546",
   "metadata": {},
   "outputs": [],
   "source": [
    "@_op.register_compute(\"simulated_quantize\")\n",
    "def simulated_quantize_compute(attrs, inputs, out_type):\n",
    "    \"\"\"模拟量化计算实现\n",
    "\n",
    "    该函数实现了模拟量化操作，包括对称量化和非对称量化，\n",
    "    并模拟了量化过程中的舍入误差。\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    attrs : Attrs\n",
    "        量化属性，包含kind、sign、rounding等参数\n",
    "    inputs : list of Tensor\n",
    "        输入张量列表，包含data、scale、clip_min、clip_max和可选的zero_point\n",
    "    out_type : Type\n",
    "        输出类型\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    list of Tensor\n",
    "        量化后并恢复的张量\n",
    "    \"\"\"\n",
    "    assert len(inputs) == 5, \"输入张量数量必须为5\"\n",
    "    assert attrs.sign, \"仅支持有符号量化\"\n",
    "    assert attrs.rounding == \"round\", \"仅支持四舍五入舍入模式\"\n",
    "\n",
    "    data = inputs[0]         # 输入数据\n",
    "    scale = inputs[1]        # 缩放因子\n",
    "    clip_min = inputs[2]     # 裁剪最小值\n",
    "    clip_max = inputs[3]     # 裁剪最大值\n",
    "    zero_point = inputs[4] if len(inputs) > 4 else None  # 零点(非对称量化使用)\n",
    "\n",
    "    # 恒等映射模式，直接返回原始数据\n",
    "    if attrs.kind == QAnnotateKind.IDENTITY:\n",
    "        return [topi.identity(data)]\n",
    "\n",
    "    # 模拟量化舍入误差\n",
    "    if zero_point is not None and attrs.kind in [QAnnotateKind.ACTIVATION, QAnnotateKind.INPUT]:\n",
    "        # 非对称量化: (x/s + zp) -> 量化 -> (q - zp) * s\n",
    "        scaled_data = topi.divide(data, scale)          # 数据除以缩放因子\n",
    "        shifted_data = topi.add(scaled_data, zero_point) # 添加零点偏移\n",
    "        clipped_data = topi.maximum(topi.minimum(shifted_data, clip_max), clip_min)  # 裁剪到有效范围\n",
    "        round_data = topi.round(clipped_data)           # 四舍五入量化\n",
    "        zero_shifted_data = topi.subtract(round_data, zero_point)  # 减去零点偏移\n",
    "        rdata = topi.multiply(zero_shifted_data, scale)  # 乘以缩放因子恢复\n",
    "    else:\n",
    "        # 对称量化: x/s -> 量化 -> q * s\n",
    "        scaled_data = topi.divide(data, scale)          # 数据除以缩放因子\n",
    "        clipped_data = topi.maximum(topi.minimum(scaled_data, clip_max), clip_min)  # 裁剪到有效范围\n",
    "        round_data = topi.round(clipped_data)           # 四舍五入量化\n",
    "        rdata = topi.multiply(round_data, scale)         # 乘以缩放因子恢复\n",
    "    \n",
    "    return [rdata]\n",
    "\n",
    "\n",
    "_reg.register_injective_schedule(\"simulated_quantize\")\n",
    "_reg.register_pattern(\"simulated_quantize\", _reg.OpPattern.ELEMWISE)\n",
    "\n",
    "def simulated_quantize(\n",
    "    data,\n",
    "    dom_scale, \n",
    "    zero_point, \n",
    "    clip_min, \n",
    "    clip_max,):\n",
    "    return _expr.Call(\n",
    "        _op.get(\"simulated_quantize\"), dom_scale, \n",
    "        zero_point, \n",
    "        clip_min, \n",
    "        clip_max, \n",
    "        # kind\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "73d80729",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %0 = relay.op.annotation.simulated_quantize(%x, %dom_scale, %zero_point, %clip_min, %clip_max, kind=1);\n",
    "#   %1 = relay.op.annotation.simulated_quantize(meta[relay.Constant][0] /* ty=Tensor[(16, 3, 3, 3), float32] span=node_conv2d.conv.weight:0:0 */, %dom_scale1, %zero_point1, %clip_min1, %clip_max1, kind=2, mode=\"i8_channel\");\n",
    "#   %2 = nn.conv2d(%0, %1, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]);\n",
    "#   %3 = relay.op.annotation.simulated_quantize(meta[relay.Constant][1] /* ty=Tensor[(16, 1, 1), float32] */, %dom_scale2, %zero_point2, %clip_min2, %clip_max2, kind=2);\n",
    "#   %4 = add(%2, %3);\n",
    "#   %5 = relay.op.annotation.simulated_quantize(%4, %dom_scale3, %zero_point3, %clip_min3, %clip_max3, kind=1);\n",
    "#   %6 = annotation.cast_hint(%5, dtype=\"int8\");\n",
    "#   annotation.stop_fusion(%6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ef72a2cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm import relay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8154447d",
   "metadata": {},
   "outputs": [
    {
     "ename": "TVMError",
     "evalue": "Traceback (most recent call last):\n  3: _ZN3tvm7runtime13PackedFun\n  2: tvm::runtime::TypedPackedFunc<tvm::relay::Call (tvm::RelayExpr, tvm::runtime::Array<tvm::RelayExpr, void>, tvm::Attrs, tvm::runtime::Array<tvm::Type, void>, tvm::Span)>::AssignTypedLambda<tvm::relay::__mk_TVM17::{lambda(tvm::RelayExpr, tvm::runtime::Array<tvm::RelayExpr, void>, tvm::Attrs, tvm::runtime::Array<tvm::Type, void>, tvm::Span)#1}>(tvm::relay::__mk_TVM17::{lambda(tvm::RelayExpr, tvm::runtime::Array<tvm::RelayExpr, void>, tvm::Attrs, tvm::runtime::Array<tvm::Type, void>, tvm::Span)#1}, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*) const\n  1: tvm::runtime::TVMMovableArgValueWithContext_::operator tvm::Span<tvm::Span>() const\n  0: _ZN3tvm7runtime6deta\n  4: _ZN3tvm7runtime13PackedFun\n  3: tvm::runtime::TypedPackedFunc<tvm::relay::Call (tvm::RelayExpr, tvm::runtime::Array<tvm::RelayExpr, void>, tvm::Attrs, tvm::runtime::Array<tvm::Type, void>, tvm::Span)>::AssignTypedLambda<tvm::relay::__mk_TVM17::{lambda(tvm::RelayExpr, tvm::runtime::Array<tvm::RelayExpr, void>, tvm::Attrs, tvm::runtime::Array<tvm::Type, void>, tvm::Span)#1}>(tvm::relay::__mk_TVM17::{lambda(tvm::RelayExpr, tvm::runtime::Array<tvm::RelayExpr, void>, tvm::Attrs, tvm::runtime::Array<tvm::Type, void>, tvm::Span)#1}, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*) const\n  2: tvm::runtime::TVMMovableArgValueWithContext_::operator tvm::Span<tvm::Span>() const\n  1: tvm::Span tvm::runtime::TVMPODValue_CRTP_<tvm::runtime::TVMArgValue>::AsObjectRef<tvm::Span>() const\n  0: _ZN3tvm7runtime6deta\n  File \"/media/pc/data/board/arria10/lxw/tasks/tvm/include/tvm/runtime/packed_func.h\", line 924\nTVMError: In function relay.ir.Call(0: RelayExpr, 1: Array<RelayExpr>, 2: Attrs, 3: Array<Type>, 4: Span) -> relay.Call: error while converting argument 4: [16:54:44] /media/pc/data/board/arria10/lxw/tasks/tvm/include/tvm/runtime/packed_func.h:2282: InternalError: Check failed: (!checked_type.defined()) is false: Expected Span, but got relay.Var\n",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mTVMError\u001b[39m                                  Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[6]\u001b[39m\u001b[32m, line 7\u001b[39m\n\u001b[32m      5\u001b[39m clip_min = relay.var(\u001b[33m\"\u001b[39m\u001b[33mclip_min\u001b[39m\u001b[33m\"\u001b[39m, shape=(), dtype=\u001b[33m\"\u001b[39m\u001b[33mfloat32\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m      6\u001b[39m clip_max = relay.var(\u001b[33m\"\u001b[39m\u001b[33mclip_max\u001b[39m\u001b[33m\"\u001b[39m, shape=(), dtype=\u001b[33m\"\u001b[39m\u001b[33mfloat32\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m----> \u001b[39m\u001b[32m7\u001b[39m x = \u001b[43msimulated_quantize\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdom_scale\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mzero_point\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mclip_min\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mclip_max\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[3]\u001b[39m\u001b[32m, line 64\u001b[39m, in \u001b[36msimulated_quantize\u001b[39m\u001b[34m(data, dom_scale, zero_point, clip_min, clip_max)\u001b[39m\n\u001b[32m     58\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34msimulated_quantize\u001b[39m(\n\u001b[32m     59\u001b[39m     data,\n\u001b[32m     60\u001b[39m     dom_scale, \n\u001b[32m     61\u001b[39m     zero_point, \n\u001b[32m     62\u001b[39m     clip_min, \n\u001b[32m     63\u001b[39m     clip_max,):\n\u001b[32m---> \u001b[39m\u001b[32m64\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_expr\u001b[49m\u001b[43m.\u001b[49m\u001b[43mCall\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m     65\u001b[39m \u001b[43m        \u001b[49m\u001b[43m_op\u001b[49m\u001b[43m.\u001b[49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43msimulated_quantize\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdom_scale\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[32m     66\u001b[39m \u001b[43m        \u001b[49m\u001b[43mzero_point\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[32m     67\u001b[39m \u001b[43m        \u001b[49m\u001b[43mclip_min\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[32m     68\u001b[39m \u001b[43m        \u001b[49m\u001b[43mclip_max\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[32m     69\u001b[39m \u001b[43m        \u001b[49m\u001b[38;5;66;43;03m# kind\u001b[39;49;00m\n\u001b[32m     70\u001b[39m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/media/pc/data/board/arria10/lxw/tasks/tvm/python/tvm/relay/expr.py:333\u001b[39m, in \u001b[36mCall.__init__\u001b[39m\u001b[34m(self, op, args, attrs, type_args, span)\u001b[39m\n\u001b[32m    331\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m type_args:\n\u001b[32m    332\u001b[39m     type_args = []\n\u001b[32m--> \u001b[39m\u001b[32m333\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m__init_handle_by_constructor__\u001b[49m\u001b[43m(\u001b[49m\u001b[43m_ffi_api\u001b[49m\u001b[43m.\u001b[49m\u001b[43mCall\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mop\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mattrs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtype_args\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mspan\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/media/pc/data/board/arria10/lxw/tasks/tvm/python/tvm/_ffi/_ctypes/object.py:167\u001b[39m, in \u001b[36mObjectBase.__init_handle_by_constructor__\u001b[39m\u001b[34m(self, fconstructor, *args)\u001b[39m\n\u001b[32m    164\u001b[39m \u001b[38;5;66;03m# assign handle first to avoid error raising\u001b[39;00m\n\u001b[32m    165\u001b[39m \u001b[38;5;66;03m# pylint: disable=not-callable\u001b[39;00m\n\u001b[32m    166\u001b[39m \u001b[38;5;28mself\u001b[39m.handle = \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m167\u001b[39m handle = \u001b[43m__init_by_constructor__\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfconstructor\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    168\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(handle, ObjectHandle):\n\u001b[32m    169\u001b[39m     handle = ObjectHandle(handle)\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/media/pc/data/board/arria10/lxw/tasks/tvm/python/tvm/_ffi/_ctypes/packed_func.py:268\u001b[39m, in \u001b[36m__init_handle_by_constructor__\u001b[39m\u001b[34m(fconstructor, args)\u001b[39m\n\u001b[32m    256\u001b[39m ret_tcode = ctypes.c_int()\n\u001b[32m    257\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[32m    258\u001b[39m     _LIB.TVMFuncCall(\n\u001b[32m    259\u001b[39m         fconstructor.handle,\n\u001b[32m   (...)\u001b[39m\u001b[32m    266\u001b[39m     != \u001b[32m0\u001b[39m\n\u001b[32m    267\u001b[39m ):\n\u001b[32m--> \u001b[39m\u001b[32m268\u001b[39m     \u001b[43mraise_last_ffi_error\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    269\u001b[39m _ = temp_args\n\u001b[32m    270\u001b[39m _ = args\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/media/pc/data/board/arria10/lxw/tasks/tvm/python/tvm/_ffi/base.py:466\u001b[39m, in \u001b[36mraise_last_ffi_error\u001b[39m\u001b[34m()\u001b[39m\n\u001b[32m    460\u001b[39m \u001b[38;5;66;03m# The exception PyObject may contain a large amount of state,\u001b[39;00m\n\u001b[32m    461\u001b[39m \u001b[38;5;66;03m# including all stack frames that may be inspected in a later\u001b[39;00m\n\u001b[32m    462\u001b[39m \u001b[38;5;66;03m# PDB post-mortem.  Therefore, we must make sure to remove the\u001b[39;00m\n\u001b[32m    463\u001b[39m \u001b[38;5;66;03m# underlying PyObject* from the C++ side after we retrieve it.\u001b[39;00m\n\u001b[32m    464\u001b[39m _LIB.TVMDropLastPythonError()\n\u001b[32m--> \u001b[39m\u001b[32m466\u001b[39m \u001b[38;5;28;01mraise\u001b[39;00m py_err\n",
      "\u001b[31mTVMError\u001b[39m: Traceback (most recent call last):\n  3: _ZN3tvm7runtime13PackedFun\n  2: tvm::runtime::TypedPackedFunc<tvm::relay::Call (tvm::RelayExpr, tvm::runtime::Array<tvm::RelayExpr, void>, tvm::Attrs, tvm::runtime::Array<tvm::Type, void>, tvm::Span)>::AssignTypedLambda<tvm::relay::__mk_TVM17::{lambda(tvm::RelayExpr, tvm::runtime::Array<tvm::RelayExpr, void>, tvm::Attrs, tvm::runtime::Array<tvm::Type, void>, tvm::Span)#1}>(tvm::relay::__mk_TVM17::{lambda(tvm::RelayExpr, tvm::runtime::Array<tvm::RelayExpr, void>, tvm::Attrs, tvm::runtime::Array<tvm::Type, void>, tvm::Span)#1}, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*) const\n  1: tvm::runtime::TVMMovableArgValueWithContext_::operator tvm::Span<tvm::Span>() const\n  0: _ZN3tvm7runtime6deta\n  4: _ZN3tvm7runtime13PackedFun\n  3: tvm::runtime::TypedPackedFunc<tvm::relay::Call (tvm::RelayExpr, tvm::runtime::Array<tvm::RelayExpr, void>, tvm::Attrs, tvm::runtime::Array<tvm::Type, void>, tvm::Span)>::AssignTypedLambda<tvm::relay::__mk_TVM17::{lambda(tvm::RelayExpr, tvm::runtime::Array<tvm::RelayExpr, void>, tvm::Attrs, tvm::runtime::Array<tvm::Type, void>, tvm::Span)#1}>(tvm::relay::__mk_TVM17::{lambda(tvm::RelayExpr, tvm::runtime::Array<tvm::RelayExpr, void>, tvm::Attrs, tvm::runtime::Array<tvm::Type, void>, tvm::Span)#1}, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*) const\n  2: tvm::runtime::TVMMovableArgValueWithContext_::operator tvm::Span<tvm::Span>() const\n  1: tvm::Span tvm::runtime::TVMPODValue_CRTP_<tvm::runtime::TVMArgValue>::AsObjectRef<tvm::Span>() const\n  0: _ZN3tvm7runtime6deta\n  File \"/media/pc/data/board/arria10/lxw/tasks/tvm/include/tvm/runtime/packed_func.h\", line 924\nTVMError: In function relay.ir.Call(0: RelayExpr, 1: Array<RelayExpr>, 2: Attrs, 3: Array<Type>, 4: Span) -> relay.Call: error while converting argument 4: [16:54:44] /media/pc/data/board/arria10/lxw/tasks/tvm/include/tvm/runtime/packed_func.h:2282: InternalError: Check failed: (!checked_type.defined()) is false: Expected Span, but got relay.Var\n"
     ]
    }
   ],
   "source": [
    "x = relay.var(\"x\", shape=(1, 3, 8, 8), dtype=\"float32\")\n",
    "# dom_scale: float32, %zero_point: int32, %clip_min: float32, %clip_max: float32,\n",
    "dom_scale = relay.var(\"dom_scale\", shape=(), dtype=\"float32\")\n",
    "zero_point = relay.var(\"zero_point\", shape=(), dtype=\"int32\")\n",
    "clip_min = relay.var(\"clip_min\", shape=(), dtype=\"float32\")\n",
    "clip_max = relay.var(\"clip_max\", shape=(), dtype=\"float32\")\n",
    "x = simulated_quantize(x, dom_scale, zero_point, clip_min, clip_max)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4ac0492",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa1b01b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "RewriteDataflowReshape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aix",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
